from numpy.linalg import norm
import numpy as np
import numba

def cosine(x, y,):
    return np.dot(x,y)/(norm(x)*norm(y))

def consine_similarity(x,y):
    '''

    :param x: (batch_size or 1,dim)
    :param y: (batch_size or 1,dim)
    :return:
    '''
    cos = np.sum(x * y,axis=-1)
    return cos/norm(x,axis=-1)/norm(y,axis=-1)

def batch_cos(arr1,arr2):
    norm1 = np.linalg.norm(arr1, axis=-1, keepdims=True)
    norm2 = np.linalg.norm(arr2, axis=-1, keepdims=True)

    arr1_norm = arr1 / norm1
    arr2_norm = arr2 / norm2

    cos = np.dot(arr1_norm, arr2_norm.T)
    return cos



def caculary_heap_index_from_features(features, thredhold=0.8):
    '''
    计算features里面相邻向量的相似度，把大于特定阈值的featyre下标放在一起
    :param features:
    :return:
    '''
    heapRelativeIndexs = []
    current_index = 0
    end_index = len(features)
    tempHeaprelativeIndex = [0]
    while current_index < end_index - 1:
        dist = cosine(features[current_index], features[current_index + 1])
        if dist >= thredhold:
            tempHeaprelativeIndex.append(current_index + 1)
        else:
            heapRelativeIndexs.append(tempHeaprelativeIndex)
            tempHeaprelativeIndex = [current_index + 1]
        current_index += 1
    heapRelativeIndexs.append(tempHeaprelativeIndex)

    heapRelativeFeatures = []
    for heap_relative_index in  heapRelativeIndexs:
        heapRelativeFeatures.append(features[heap_relative_index[0]:heap_relative_index[-1] + 1])
    return heapRelativeIndexs,heapRelativeFeatures



def caculate_heaps_similary_score(heap_features_A,heap_features_B):
    '''
    heap_features_A 和 heap_features_B里面所有的feature的平均相似度
    :param feature:
    :param features_heap:
    :return:
    '''
    similary_score = []
    for feature in heap_features_A:
        feature = np.array(feature)
        feature = feature.reshape((1, -1))
        single_result = consine_similarity(feature, heap_features_B)
        similary_score.extend(single_result)
    return sum(similary_score)/len(similary_score)


def caculate_heaps_similary_dict(heapRelativeIndexs_A,heapRelativeFeatures_A,heapRelativeIndexs_,heaps_B_features):
    '''
    计算
    :param heaps_A_indexs:
    :param heaps_A_features:
    :param heaps_B_indexs:
    :param heaps_B_features:
    :return:
    '''
    heap_dict = {}
    for index, heaps_A_feature in enumerate(heaps_A_features):
        total_scores = []
        jieshuo_features = jieshuo_image_features[jieshuo_heap[0]:jieshuo_heap[-1] + 1]
        for origin_heap in origin_heaps:
            origin_features = origin_image_features[origin_heap[0]:origin_heap[-1] + 1]
            similary_score = 0
            if len(origin_heap) >= len(jieshuo_heap):
                similary_score = caculate_heap_and_heap_similary_dict(jieshuo_features, origin_features)
            total_scores.append(similary_score)
        heap_dict[tuple(jieshuo_heap)] = sorted(range(len(total_scores)), key=lambda i: total_scores[i], reverse=True)[
                                         :10]
    return heap_dict

def get_heap_like_index(jieshuo_heaps,origin_heaps,jieshuo_image_features,origin_image_features):
    heap_dict = {}
    for jieshuo_index, jieshuo_heap  in enumerate(jieshuo_heaps):
        total_scores = []
        jieshuo_features = jieshuo_image_features[jieshuo_heap[0]:jieshuo_heap[-1] +1]
        for  origin_heap in origin_heaps:
            origin_features = origin_image_features[origin_heap[0]:origin_heap[-1] +1]
            similary_score = 0
            if len(origin_heap) >= len(jieshuo_heap):
                similary_score = caculate_heap_and_heap_similary_dict(jieshuo_features,origin_features)
            total_scores.append(similary_score)
        heap_dict[tuple(jieshuo_heap)] = sorted(range(len(total_scores)), key=lambda i: total_scores[i],reverse=True)[:10]
    return heap_dict

if __name__ == '__main__':
    a = np.arange(2000 * 3000).reshape((2000,3000))
    b = np.arange(20000  * 3000).reshape((20000,3000))
    print(batch_cos(a,b))


